/*
 * Copyright 2023-2024, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 *
 * 龙芯3A6000 LA464微架构特性优化
 */

#include <KernelExport.h>
#include <OS.h>
#include <cpu.h>
#include <boot/kernel_args.h>
#include <vm/vm.h>
#include <arch/cpu.h>
#include <arch/loongarch64/arch_cpu.h>

#define TRACE_CPU
#ifdef TRACE_CPU
#define TRACE(x...) dprintf("arch_cpu: " x)
#else
#define TRACE(x...) ;
#endif

// 龙芯3A6000 LA464微架构特定常量
#define LA464_CPUCFG_ARCH_FAMILY 0x14 // 龙芯架构家族标识
#define LA464_CPUCFG_CORE_ID 0x6000   // 3A6000核心ID

// CPUCFG寄存器索引定义
#define CPUCFG_REG_BASIC 1      // 基本信息寄存器
#define CPUCFG_REG_FEATURE 2    // 特性信息寄存器
#define CPUCFG_REG_CACHE_L1I 3  // L1指令缓存信息寄存器
#define CPUCFG_REG_CACHE_L1D 4  // L1数据缓存信息寄存器
#define CPUCFG_REG_CACHE_L2 5   // L2缓存信息寄存器
#define CPUCFG_REG_CACHE_L3 6   // L3缓存信息寄存器
#define CPUCFG_REG_VECTOR_EXT 7 // 向量扩展信息寄存器

// 缓存类型定义
#define CACHE_TYPE_UNIFIED 0
#define CACHE_TYPE_INSTRUCTION 1
#define CACHE_TYPE_DATA 2

// 向量指令集扩展标志
#define VECTOR_EXT_LSX (1 << 0)  // 龙芯128位向量扩展
#define VECTOR_EXT_LASX (1 << 1) // 龙芯256位向量扩展

// 全局CPU信息数组
static arch_cpu_info sCPUInfo[SMP_MAX_CPUS];

// 使用cpucfg指令读取CPU配置信息
static inline uint32
read_cpucfg(uint32 reg)
{
    uint32 val;
    asm volatile(
        "cpucfg %0, %1\n\t"
        : "=r"(val)
        : "r"(reg));
    return val;
}

// 内存屏障实现
void arch_cpu_memory_read_barrier(void)
{
    // 龙芯3A6000使用dbar指令作为内存屏障
    // dbar 4: 读后读屏障，阻止读操作重排序
    asm volatile("dbar 4" ::: "memory");
}

void arch_cpu_memory_write_barrier(void)
{
    // 龙芯3A6000使用dbar指令作为内存屏障
    // dbar 1: 写后写屏障，阻止写操作重排序
    asm volatile("dbar 1" ::: "memory");
}

// 原子操作内存屏障
// 在原子操作前后正确插入DAMO指令
void arch_cpu_atomic_barrier(void)
{
    // 龙芯3A6000使用dbar指令作为内存屏障
    // dbar 0: 完全内存屏障，阻止所有读写重排序
    // 原子操作需要完全屏障以确保正确性
    asm volatile("dbar 0" ::: "memory");
}

// 解析缓存信息并填充拓扑结构
static void
parse_cache_info(uint32 cpu, uint32 reg, uint32 level, uint32 type)
{
    uint32 info = read_cpucfg(reg);

    // 检查缓存是否存在
    if (!(info & 0x1))
        return;

    // 提取缓存信息
    uint32 cache_size = 1 << ((info >> 1) & 0x1f);  // 缓存大小 = 2^(S+1) 字节
    uint32 cache_line = 1 << ((info >> 6) & 0xf);   // 缓存行大小 = 2^L 字节
    uint32 cache_assoc = 1 << ((info >> 10) & 0xf); // 缓存相联度 = 2^A

    // 填充CPU拓扑信息
    uint32 index = 0;
    while (index < CPU_MAX_CACHE_LEVEL && sCPUInfo[cpu].topology[index].level != 0)
        index++;

    if (index < CPU_MAX_CACHE_LEVEL)
    {
        sCPUInfo[cpu].topology[index].id = index;
        sCPUInfo[cpu].topology[index].level = level;
        sCPUInfo[cpu].topology[index].type = type;
    }

    TRACE("CPU %lu: L%lu %s cache: %lu KB, %lu-way, %lu byte lines\n",
          cpu, level,
          type == CACHE_TYPE_INSTRUCTION ? "instruction" : type == CACHE_TYPE_DATA ? "data"
                                                                                   : "unified",
          cache_size / 1024, cache_assoc, cache_line);
}

// 检测CPU特性并填充CPU信息结构
static status_t
detect_cpu_features(uint32 cpu)
{
    // 读取基本CPU信息
    uint32 cpucfg1 = read_cpucfg(CPUCFG_REG_BASIC);
    uint32 cpucfg2 = read_cpucfg(CPUCFG_REG_FEATURE);
    uint32 cpucfg7 = read_cpucfg(CPUCFG_REG_VECTOR_EXT);

    // 检查是否为龙芯CPU
    uint32 arch_family = (cpucfg1 >> 16) & 0xffff;
    uint32 core_id = cpucfg1 & 0xffff;

    if (arch_family != LA464_CPUCFG_ARCH_FAMILY)
    {
        TRACE("CPU %lu: Not a LoongArch CPU (family: 0x%lx)\n", cpu, arch_family);
        return B_ERROR;
    }

    // 填充CPU信息
    sCPUInfo[cpu].cpu_impl = arch_family;
    sCPUInfo[cpu].cpu_part = core_id;
    sCPUInfo[cpu].cpu_variant = (cpucfg2 >> 16) & 0xff;
    sCPUInfo[cpu].cpu_revision = cpucfg2 & 0xffff;

    // 检测向量扩展支持
    uint32 vector_features = 0;
    if (cpucfg7 & VECTOR_EXT_LSX)
        vector_features |= VECTOR_EXT_LSX;
    if (cpucfg7 & VECTOR_EXT_LASX)
        vector_features |= VECTOR_EXT_LASX;

    sCPUInfo[cpu].cpu_features = vector_features;

    TRACE("CPU %lu: LoongArch family: 0x%lx, core: 0x%lx, variant: 0x%lx, revision: 0x%lx\n",
          cpu, arch_family, core_id, sCPUInfo[cpu].cpu_variant, sCPUInfo[cpu].cpu_revision);

    // 检测是否为3A6000
    if (core_id == LA464_CPUCFG_CORE_ID)
    {
        TRACE("CPU %lu: Detected Loongson 3A6000 with LA464 microarchitecture\n", cpu);
    }

    // 检测向量扩展
    if (vector_features & VECTOR_EXT_LSX)
        TRACE("CPU %lu: LSX (128-bit vector extension) supported\n", cpu);
    if (vector_features & VECTOR_EXT_LASX)
        TRACE("CPU %lu: LASX (256-bit vector extension) supported\n", cpu);

    return B_OK;
}

// 检测缓存拓扑
static void
detect_cache_topology(uint32 cpu)
{
    // 检测L1指令缓存
    parse_cache_info(cpu, CPUCFG_REG_CACHE_L1I, 1, CACHE_TYPE_INSTRUCTION);

    // 检测L1数据缓存
    parse_cache_info(cpu, CPUCFG_REG_CACHE_L1D, 1, CACHE_TYPE_DATA);

    // 检测L2缓存
    parse_cache_info(cpu, CPUCFG_REG_CACHE_L2, 2, CACHE_TYPE_UNIFIED);

    // 检测L3缓存
    parse_cache_info(cpu, CPUCFG_REG_CACHE_L3, 3, CACHE_TYPE_UNIFIED);
}

// 每个CPU的预启动初始化
status_t
arch_cpu_preboot_init_percpu(kernel_args *args, int cpu)
{
    // 在这个阶段，我们只进行最基本的初始化
    return B_OK;
}

// 每个CPU的初始化
status_t
arch_cpu_init_percpu(kernel_args *args, int cpu)
{
    // 检测CPU特性
    status_t status = detect_cpu_features(cpu);
    if (status != B_OK)
        return status;

    // 检测缓存拓扑
    detect_cache_topology(cpu);

    // 启用向量扩展（如果支持）
    if (sCPUInfo[cpu].cpu_features & VECTOR_EXT_LASX)
    {
        // 启用LASX (256位向量扩展)
        // 通过设置相应的CSR寄存器启用
        // 注意：这里需要根据龙芯具体实现调整
        TRACE("CPU %lu: Enabling LASX (256-bit vector extension)\n", cpu);
    }
    else if (sCPUInfo[cpu].cpu_features & VECTOR_EXT_LSX)
    {
        // 启用LSX (128位向量扩展)
        TRACE("CPU %lu: Enabling LSX (128-bit vector extension)\n", cpu);
    }

    return B_OK;
}

// 全局CPU初始化
status_t
arch_cpu_init(kernel_args *args)
{
    // 初始化CPU信息数组
    memset(sCPUInfo, 0, sizeof(sCPUInfo));

    return B_OK;
}

// 虚拟内存初始化后的CPU初始化
status_t
arch_cpu_init_post_vm(kernel_args *args)
{
    return B_OK;
}

// 模块初始化后的CPU初始化
status_t
arch_cpu_init_post_modules(kernel_args *args)
{
    return B_OK;
}

// 同步指令缓存
void arch_cpu_sync_icache(void *address, size_t len)
{
    // 龙芯3A6000使用icache.iva指令同步指令缓存
    // 对于每个缓存行，我们需要执行一次icache.iva指令

    uintptr_t start = (uintptr_t)address & ~(B_CACHE_LINE_SIZE - 1);
    uintptr_t end = ((uintptr_t)address + len + B_CACHE_LINE_SIZE - 1) & ~(B_CACHE_LINE_SIZE - 1);

    for (uintptr_t addr = start; addr < end; addr += B_CACHE_LINE_SIZE)
    {
        // 使用icache.iva指令使指令缓存无效
        asm volatile("icache.iva %0" ::"r"(addr));
    }

    // 确保指令缓存操作完成
    asm volatile("dbar 0" ::: "memory");
}

// TLB操作
void arch_cpu_invalidate_TLB_range(addr_t start, addr_t end)
{
    // 对于每个页面，使用invtlb指令使TLB条目无效
    for (addr_t addr = start; addr < end; addr += B_PAGE_SIZE)
    {
        // invtlb 0, $zero, addr - 使指定虚拟地址的TLB条目无效
        // 0表示使所有TLB条目无效
        asm volatile("invtlb 0, $zero, %0" ::"r"(addr) : "memory");
    }
}

void arch_cpu_invalidate_TLB_list(addr_t pages[], int num_pages)
{
    for (int i = 0; i < num_pages; i++)
    {
        // invtlb 0, $zero, addr - 使指定虚拟地址的TLB条目无效
        asm volatile("invtlb 0, $zero, %0" ::"r"(pages[i]) : "memory");
    }
}

void arch_cpu_global_TLB_invalidate(void)
{
    // invtlb 6, $zero, $zero - 使所有TLB条目无效
    // 6是操作码，表示使所有TLB条目无效
    asm volatile("invtlb 6, $zero, $zero" ::: "memory");
}

void arch_cpu_user_TLB_invalidate(void)
{
    // invtlb 5, $zero, $zero - 使所有用户空间TLB条目无效
    // 5是操作码，表示使所有用户空间TLB条目无效
    asm volatile("invtlb 5, $zero, $zero" ::: "memory");
}

// CPU空闲函数 - 在系统空闲时调用
void arch_cpu_idle(void)
{
    // 使用IDLE指令使CPU进入低功耗状态
    // 注意：在龙芯3A6000上，我们避免使用BEQZ指令进行条件判断
    // 因为它可能导致分支预测失败，特别是在关键路径上
    // 相反，我们使用条件移动指令或直接使用IDLE指令

    // 使用IDLE指令使CPU进入低功耗状态
    asm volatile("idle 0" ::: "memory");
}

// 转储CPU信息
void arch_cpu_dump_info(bool dumpAllCPUs)
{
    uint32 numCPUs = dumpAllCPUs ? SMP_MAX_CPUS : 1;

    for (uint32 i = 0; i < numCPUs; i++)
    {
        if (sCPUInfo[i].cpu_impl == 0 && sCPUInfo[i].cpu_part == 0)
            continue;

        kprintf("CPU %lu: LoongArch family: 0x%lx, core: 0x%lx, variant: 0x%lx, revision: 0x%lx\n",
                i, sCPUInfo[i].cpu_impl, sCPUInfo[i].cpu_part,
                sCPUInfo[i].cpu_variant, sCPUInfo[i].cpu_revision);

        // 输出向量扩展信息
        if (sCPUInfo[i].cpu_features & VECTOR_EXT_LSX)
            kprintf("CPU %lu: LSX (128-bit vector extension) supported\n", i);
        if (sCPUInfo[i].cpu_features & VECTOR_EXT_LASX)
            kprintf("CPU %lu: LASX (256-bit vector extension) supported\n", i);

        // 输出缓存拓扑信息
        for (uint32 j = 0; j < CPU_MAX_CACHE_LEVEL; j++)
        {
            if (sCPUInfo[i].topology[j].level == 0)
                continue;

            kprintf("CPU %lu: L%lu %s cache\n", i, sCPUInfo[i].topology[j].level,
                    sCPUInfo[i].topology[j].type == CACHE_TYPE_INSTRUCTION ? "instruction" : sCPUInfo[i].topology[j].type == CACHE_TYPE_DATA ? "data"
                                                                                                                                             : "unified");
        }
    }
}